package org.zjx.server;

import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.client.SocketClient;
import org.zjx.core.SocketConfig;
import org.zjx.exception.RegistrationFailedException;
import org.zjx.exception.RpcTimeoutException;
import org.zjx.message.IMessage;
import org.zjx.message.messages.*;

import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public class MasterConnector {
    private static final Logger logger = LoggerFactory.getLogger(MasterConnector.class);
    private final ClusterServer server;
    private final String masterHost;
    private final int masterPort;
    private SocketClient masterClient;
    private final AtomicLong lastHeartbeatTime = new AtomicLong(System.currentTimeMillis());
    private ScheduledFuture<?> heartbeatFuture;
    private final AtomicBoolean isReconnecting = new AtomicBoolean(false);

    public MasterConnector(ClusterServer server, String host, int port) {
        this.server = server;
        this.masterHost = host;
        this.masterPort = port;
    }

    public void connect() {
        try {
            logger.info("Connecting to master at {}:{}", masterHost, masterPort);
            initMasterClient();
            registerWorker();
        } catch (Exception e) {
            logger.error("Connection to master failed", e);
            scheduleReconnect();
        }
    }

    public RouteResponse queryRoute(String serviceName) {
        try {
            if (masterClient == null || !masterClient.isConnected()) {
                logger.warn("Not connected to master, cannot query route");
                return RouteResponse.unavailable(serviceName, "Not connected to master");
            }

            RouteQueryMessage query = new RouteQueryMessage();
            query.setServiceName(serviceName);

            IMessage response = masterClient.sendAndWait(query, server.getConfig().getRequestTimeout());

            if (response instanceof RouteResponse) {
                return (RouteResponse) response;
            }

            logger.error("Unexpected response type for route query: {}",
                    response != null ? response.getClass().getSimpleName() : "null");
            return RouteResponse.unavailable(serviceName, "Invalid response from master");
        } catch (RpcTimeoutException e) {
            logger.error("Route query timeout: {}", serviceName, e);
            return RouteResponse.unavailable(serviceName, "Route query timeout");
        } catch (Exception e) {
            logger.error("Route query failed: {}", serviceName, e);
            return RouteResponse.unavailable(serviceName, "Route query failed");
        }
    }

    private void initMasterClient() {
        SocketConfig masterConfig = new SocketConfig()
                .setServerHost(masterHost)
                .setPort(masterPort)
                .setHeartbeatInterval(server.getConfig().getHeartbeatInterval())
                .setReconnectInterval(server.getConfig().getReconnectInterval())
                .setRequestTimeout(server.getConfig().getRequestTimeout());

        masterClient = new SocketClient(masterConfig);
        masterClient.setMessageListener(this::handleMasterMessage);
        masterClient.connect();
    }

    private void handleMasterMessage(IMessage message) {
        if (message instanceof HeartbeatResponse response) {
            // 关键修复：验证响应中的 workerId 是否匹配
            if (response.getWorkerId() != null && response.getWorkerId().equals(server.getServerId())) {
                updateLastHeartbeat();
                logger.debug("Received valid heartbeat response");
            } else {
                // 客户端心跳响应（降低日志级别）
                logger.trace("Received client heartbeat response");
            }
        } else if (message instanceof AckMessage) {
            logger.info("Received registration acknowledgement");
        }
    }

    @SneakyThrows
    private void registerWorker() {
        int retryCount = 0;
        long baseTimeout = 15000; // 初始超时时间15秒
        long maxTimeout = 60000;  // 最大超时时间60秒

        while (retryCount < 5) {
            try {
                logger.info("Registering worker (attempt {})", retryCount + 1);
                RegisterWorkerMessage regMsg = createRegistrationMessage();
                IMessage response = masterClient.sendAndWait(regMsg, baseTimeout);

                if (response instanceof AckMessage) {
                    logger.info("Worker registration acknowledged");
                    startHeartbeat();
                    return;
                }

                logger.warn("Unexpected registration response: {}",
                        response != null ? response.getClass().getSimpleName() : "null");
            } catch (Exception e) {
                logger.error("Registration attempt {} failed", retryCount + 1, e);
            }

            // 指数退避
            long delay = (long) (Math.pow(2, retryCount) * 3000);
            logger.info("Retrying registration in {}ms", delay);
            try {
                Thread.sleep(delay);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }

            // 增加超时时间，但不超过最大值
            baseTimeout = Math.min(baseTimeout + 5000, maxTimeout);
            retryCount++;
        }

        throw new RegistrationFailedException("Worker registration failed after 5 attempts");
    }

    private RegisterWorkerMessage createRegistrationMessage() {
        RegisterWorkerMessage regMsg = new RegisterWorkerMessage();
        regMsg.setWorkerId(server.getServerId());

        // 关键修复：使用Worker自身的服务地址和端口
        // 直接使用 serverHost 和 port
        String workerHost = server.getConfig().getServerHost();
        int workerPort = server.getConfig().getPort();

        regMsg.setHost(workerHost);
        regMsg.setPort(workerPort);

        regMsg.setServices(server.getServiceRegistry().getRegisteredServices());

        // 添加日志验证
        logger.info("Registering worker at {}:{} with services: {}",
                workerHost,
                workerPort,
                regMsg.getServices());

        return regMsg;
    }

    private void startHeartbeat() {
        // 确保只启动一个心跳任务
        if (heartbeatFuture != null && !heartbeatFuture.isDone()) {
            heartbeatFuture.cancel(true);
        }

        ScheduledExecutorService executor = server.getBackgroundExecutor();
        heartbeatFuture = executor.scheduleAtFixedRate(() -> {
            // 检查连接状态
            if (masterClient == null || !masterClient.isConnected()) {
                logger.warn("Not connected, skipping heartbeat");
                return;
            }

            // 检查心跳超时
            if (System.currentTimeMillis() - lastHeartbeatTime.get() > server.getConfig().getHeartbeatTimeout()) {
                logger.warn("Master heartbeat timeout, reconnecting...");
                reconnect();
                return;
            }

            try {
                // 创建并发送心跳
                HeartbeatMessage heartbeat = new HeartbeatMessage();
                heartbeat.setWorkerId(server.getServerId());
                heartbeat.setMessageId("HB_" + System.currentTimeMillis());

                logger.debug("Sending heartbeat to master: workerId={}", server.getServerId());
                masterClient.sendMessage(heartbeat);
            } catch (Exception e) {
                logger.error("Heartbeat failed", e);
                reconnect();
            }
        }, 0, server.getConfig().getHeartbeatInterval(), TimeUnit.MILLISECONDS);
    }

    void updateLastHeartbeat() {
        lastHeartbeatTime.set(System.currentTimeMillis());
    }

    private void reconnect() {
// 添加重入锁，防止并发重连
        if (isReconnecting.get()) return;
        isReconnecting.set(true);

        try {
            logger.info("Reconnecting to master...");

            // 1. 停止当前心跳任务
            if (heartbeatFuture != null) {
                heartbeatFuture.cancel(true);
                heartbeatFuture = null;
            }

            // 2. 断开旧连接
            if (masterClient != null) {
                masterClient.disconnect();
                masterClient = null;
            }

            // 3. 建立新连接
            initMasterClient();

            // 4. 重新注册
            registerWorker();

            // 5. 重启心跳
            startHeartbeat();
        } catch (Exception e) {
            logger.error("Reconnect failed", e);

            // 指数退避重试
            long delay = Math.min(server.getConfig().getReconnectInterval() * 2, 60000);
            logger.info("Scheduling reconnect in {}ms", delay);

            server.getBackgroundExecutor().schedule(() -> {
                isReconnecting.set(false);
                reconnect();
            }, delay, TimeUnit.MILLISECONDS);
        } finally {
            isReconnecting.set(false);
        }
    }

    private void scheduleReconnect() {
        server.getBackgroundExecutor().schedule(this::reconnect,
                server.getConfig().getReconnectInterval(), TimeUnit.MILLISECONDS);
    }

    public void disconnect() {
        // 1. 取消心跳任务
        if (heartbeatFuture != null) {
            heartbeatFuture.cancel(true);
            heartbeatFuture = null;
        }

        // 2. 断开与主服务器的连接
        if (masterClient != null) {
            masterClient.disconnect();
            masterClient = null;
        }

        logger.info("Master connector disconnected");
    }
}
